import gym
from gym import spaces
import numpy as np
import pandas as pd
import yfinance as yf
import matplotlib.pyplot as plt
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
# Download data
df = df
df = df[['Open', 'High', 'Low', 'Close', 'Volume']]
df['Returns'] = df['Close'].pct_change()
df.dropna(inplace=True)
# Custom Trading Environment
class TradingEnv(gym.Env):
def __init__(self, df):
super(TradingEnv, self).__init__()
self.df = df.reset_index()
self.max_steps = len(df) - 1
self.current_step = 0
self.initial_balance = 10000
self.balance = self.initial_balance
self.shares_held = 0
self.net_worth = self.initial_balance
# Actions: 0 = Hold, 1 = Buy, 2 = Sell
self.action_space = spaces.Discrete(3)
# Observation: [Open, High, Low, Close, Volume, Shares Held, Balance]
self.observation_space = spaces.Box(
low=0, high=np.inf, shape=(7,), dtype=np.float32
)
def _get_obs(self):
obs = np.array([
self.df.loc[self.current_step, 'Open'],
self.df.loc[self.current_step, 'High'],
self.df.loc[self.current_step, 'Low'],
self.df.loc[self.current_step, 'Close'],
self.df.loc[self.current_step, 'Volume'],
self.shares_held,
self.balance
], dtype=np.float32)
return obs
def step(self, action):
current_price = self.df.loc[self.current_step, 'Close']
done = False
reward = 0
# Action 1: Buy
if action == 1 and self.balance >= current_price:
shares_bought = self.balance // current_price
cost = shares_bought * current_price
self.balance -= cost
self.shares_held += shares_bought
# Action 2: Sell
elif action == 2 and self.shares_held > 0:
revenue = self.shares_held * current_price
self.balance += revenue
self.shares_held = 0
self.current_step += 1
self.net_worth = self.balance + self.shares_held * current_price
reward = self.net_worth - self.initial_balance
if self.current_step >= self.max_steps:
done = True
info = {'net_worth': self.net_worth}
return self._get_obs(), reward, done, info
def reset(self):
self.balance = self.initial_balance
self.shares_held = 0
self.net_worth = self.initial_balance
self.current_step = 0
return self._get_obs()
def render(self, mode='human'):
print(f'Step: {self.current_step}')
print(f'Balance: {self.balance}')
print(f'Shares held: {self.shares_held}')
print(f'Net worth: {self.net_worth}')
# Create environment
env = TradingEnv(df)
env = DummyVecEnv([lambda: env])
# Train PPO agent
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=20000)
# Test agent
obs = env.reset()
done = False
actions = []
prices = df['Close'].values
while not done:
action, _states = model.predict(obs)
obs, reward, done, info = env.step(action)
actions.append(action[0])
# Visualize results with Buy/Sell markers
plt.figure(figsize=(15, 7))
plt.plot(df.index, df['Close'], label='Close Price')
buy_signals = [i for i, a in enumerate(actions) if a == 1]
sell_signals = [i for i, a in enumerate(actions) if a == 2]
plt.scatter(df.index[buy_signals], df['Close'].iloc[buy_signals], marker='^', color='g', label='Buy Signal', s=100)
plt.scatter(df.index[sell_signals], df['Close'].iloc[sell_signals], marker='v', color='r', label='Sell Signal', s=100)
plt.title('AAPL Close Price and RL Trading Signals')
plt.legend()
plt.show()
